from audioop import reverse
import numpy as np
import random
import copy

class Particle:

    def __init__(self, obj_index, inner_id, config):
        """_summary_
        Class Particle implement the single particle (in one swarm for objective i) actions.
        Modify the solution arch.
        Args:
            obj_index (int): the i-th objective for the i-th swarm mainly concerned.
            inner_id (int): the j-th particle in the i-th swarm.
            config (dict): include the information about the decision space dimesion, objective space dimesion and so on.
        """

        self.obj_index = obj_index
        self.inner_id = inner_id
        self.config = config
        
        ## solution
        self.eps = self.config['UpperEpsilon'] * np.random.rand(1)
        self.x = self.config['LowerBound'] + (self.config['UpperBound'] - self.config['LowerBound']) * np.random.rand(self.config['DsDim'])
        self.x = np.concatenate([self.eps, self.x])

        ## velocity
        self.eps_v = -self.config['UpperEpsVelocity'] + 2 * self.config['UpperEpsVelocity'] * np.random.rand(1)
        self.v = self.config['LowerVelocity'] + (self.config['UpperVelocity'] - self.config['LowerVelocity']) * np.random.rand(self.config['DsDim'])
        self.v = np.concatenate([self.eps_v, self.v])
        
        ## fitness
        self.fitness = np.ones(self.config['OsDim'])
        self.eval()
        
        self.distance = 0.0  ## for crowding ditance sort

    def update_particle(self, archive, pBest, gBest):
        """_summary_

        Args:
            archive (Dict): The external archive for information sharing, using two elements, the 'size' and the 'set'. 
        """
        if archive['size'] > 0:
            ai = random.randint(0, archive['size'] - 1)
        else:
            ai = random.randint(0, self.config['OsDim'] - 1)
            while ai == self.obj_index:
                ai = random.randint(0, self.config['OsDim'] - 1)
            if ai == self.obj_index or ai >= self.config['OsDim'] or ai < 0:
                raise ValueError
        
        r1 = np.random.rand(self.x.shape[0])
        r2 = np.random.rand(self.x.shape[0])
        r3 = np.random.rand(self.x.shape[0])

        if archive['size'] > 0:
            self.v = self.config['inertiaw'] * self.v
            + self.config['c1'] * r1 * (pBest[self.obj_index][self.inner_id].x - self.x)
            + self.config['c2'] * r2 * (gBest[self.obj_index].x - self.x)
            + self.config['c3'] * r3 * (archive['set'][ai].x - self.x)
        else:
            self.v = self.config['inertiaw'] * self.v
            + self.config['c1'] * r1 * (pBest[self.obj_index][self.inner_id].x - self.x)
            + self.config['c2'] * r2 * (gBest[self.obj_index].x - self.x)
            + self.config['c3'] * r3 * (gBest[ai].x - self.x)

        ## clamp velocity into [v_lower, v_upper]
        self.v[0] = np.clip(self.v[0], a_min=-self.config['UpperEpsVelocity'], a_max=self.config['UpperEpsVelocity'])
        self.v[1:] = np.clip(self.v[1:], a_min=self.config['LowerVelocity'], a_max=self.config['UpperVelocity'])
        
        ## update solution x
        self.x += self.v
        
        ## clamp solution
        self.x[0] = np.clip(self.x[0], a_min=0., a_max=self.config['UpperEpsilon'])
        self.x[1:] = np.clip(self.x[1:], a_min=self.config['LowerBound'], a_max=self.config['UpperBound'])
        ## update fitness
        self.eval()
    
    
    def eval(self, ):
        """_summary_
        Update the particle's fitness.
        Args:
            evaluater (Evaluation): give the function eval result, which return the fitness.
        """
        res = self.config['evaluationer'].eval(self.x)

        if res.shape[0] != self.config['OsDim']:
            raise ValueError

        self.fitness = res


    def deepcopy(self, copy_object):
        """_summary_
        Let the copy_objectt(a particle)'s information be copyed to 'this' particle.

        Args:
            copy_objec (_type_): type: Particle
        """
        self.obj_index = copy_object.obj_index
        self.x = copy.deepcopy(copy_object.x)
        self.v = copy.deepcopy(copy_object.v)
        self.fitness = copy.deepcopy(copy_object.fitness)
    
    def els(self, ):
        
        # self.x += self.config['UpperBound'] * np.random.randn(self.x.shape[0])
        # self.x = np.clip(self.x, a_min=self.config['LowerBound'], a_max=self.config['UpperBound'])
        d = random.randint(1, self.config['DsDim'])
        self.x[d] += self.config['UpperBound'] * random.gauss(0., 1.)
        
        ## here, we following the offical code (C plus plus version), if the x[d] out of range, we randomly sample the x[d] into [lower, upper]
        if self.x[d] < self.config['LowerBound'] or self.x[d] > self.config['UpperBound']:
            self.x[d] = random.uniform(self.config['LowerBound'], self.config['UpperBound'])

        self.eval()

    def reset_distance(self, ):
        self.distance = 0.0

class CMPSO:

    def __init__(self, config):
        """_summary_
        The main body for CMPSO.
        the main process:
            1) maintain a archive.
            2) build a swarms, each swarm have a mainly objective and maintain some population in the swarm.
        Args:
            config (Dict): Include the information about the decision space dimesion, objective space dimesion and so on.
        """
        
        self.config = config
        self.DsDim = self.config['DsDim']
        self.OsDim = self.config['OsDim']
        
        ## Init the swarms, where each swarm mainly for one objective.
        self.swarms = []
        for obj_index in range(self.OsDim):
            #swarm for i-th objective
            obj_i_swarm = [Particle(obj_index=obj_index, inner_id=i, config=self.config) for i in range(self.config['NP'])]
            if len(obj_i_swarm) != self.config['NP']:
                raise ValueError

            self.swarms.append(obj_i_swarm)

        ## Init pBest, deepcopy from swarms firstly.
        self.pBest = []
        for obj_index in range(self.OsDim):

            # pBest for i-th objective, init pBest wit same like the init particle for coresponding position.
            obj_i_pBest = []
            for i in range(self.config['NP']):
                temp = Particle(obj_index=obj_index, inner_id=i, config=self.config)
                temp.deepcopy(self.swarms[obj_index][i])
                obj_i_pBest.append(temp)
            if len(obj_i_pBest) != self.config['NP']:
                raise ValueError

            self.pBest.append(obj_i_pBest) 

        ## Init and update gBest
        self.gBest = [Particle(obj_index=obj_index, inner_id=-1, config=self.config) for obj_index in range(self.OsDim)]
        self.update_gBest()


        ## Init the archive
        self.archive = {
            'size': 0,  ## current size
            'upper_size': self.config['arc_size'], ## the upper size of the archive
            'set': [] ## to stroe those particle
        }
        self.update_archive()

    def update_pBest(self, ):

        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                if self.swarms[obj_index][i].fitness[obj_index] < self.pBest[obj_index][i].fitness[obj_index]:
                    self.pBest[obj_index][i].deepcopy(self.swarms[obj_index][i])

    def update_gBest(self, ):

        for obj_index in range(self.OsDim):
            min_value_index = 0
            for index in range(self.config['NP']):
                if self.pBest[obj_index][index].fitness[obj_index] < self.pBest[obj_index][min_value_index].fitness[obj_index]:
                    min_value_index = index
            
            if self.pBest[obj_index][min_value_index].fitness[obj_index] < self.gBest[obj_index].fitness[obj_index]:
                self.gBest[obj_index].deepcopy(self.pBest[obj_index][min_value_index])   

    def update_archive(self, ):
        
        def dominate(fitness_1, fitness_2):
            """
            return 0/-1/1, where 0 represent co-dominate, -1 represent fit_2 domainte fit_1, 1 represent fit 1 domainate fit_2.
            """
            better = (fitness_1 < fitness_2).sum()
            same = (fitness_1 == fitness_2).sum()
            worse = (fitness_1 > fitness_2).sum()

            if same == fitness_1.shape[0]:
                return 0
            elif better == 0:
                return -1
            elif worse == 0:
                return 1
            else:
                return 0

        def nondominatedSet(S_set):
            """_summary_
                Get non-dominated Set R and return.
            Args:
                S_set (_type_): _description_
            Retrun:
                R_set (List):
            """
            size = len(S_set)
            R_set = []

            for i in range(size):
                non = True
                parttcle_a = S_set[i]
                for j in range(size):
                    if i != j:
                        parttcle_b = S_set[j]
                        k = dominate(parttcle_a.fitness, parttcle_b.fitness)
                        if k < 0:
                            non = False
                            break
                if non:
                    R_set.append(parttcle_a)
            return R_set            

        def crowding_distance_assignment(R_set):
            """_summary_
                Give each particles a crowding distance, and return the "sorted(by distance)" set.
            Args:
                R_set (List): 
            """
            ## zero the distance for particles in R_set
            size = len(R_set)
            for particle in R_set:
                particle.reset_distance()
            for sort_obj_index in range(self.OsDim):
                # sort by the i-th objective
                R_set.sort(key=lambda particle: particle.fitness[sort_obj_index])
                R_set[0].distance = 1e6
                R_set[size - 1].distance = 1e6
                fitmin = R_set[0].fitness[sort_obj_index]
                fitmax = R_set[size - 1].fitness[sort_obj_index]

                for index in range(1, size - 1):
                    R_set[index].distance += (R_set[index + 1].fitness[sort_obj_index] - R_set[index - 1].fitness[sort_obj_index]) / (fitmax - fitmin)
            
            # sort by distance
            R_set.sort(key=lambda particle: particle.distance, reverse=True)

            return R_set

        S = [] 
        
        ## Add pBest of each particle in each swarm into set s
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                S.append(self.pBest[obj_index][i])


        ## Add all the solution in the old archive into the set S
        for i in range(self.archive['size']):
            S.append(self.archive['set'][i])
        

        ## ELS strategy
        for i in range(self.archive['size']):
            temp = Particle(obj_index=self.archive['set'][i].obj_index, inner_id=-1, config=self.config)
            temp.deepcopy(self.archive['set'][i])
            temp.els()
            S.append(temp)
        

        R = nondominatedSet(S)


        if len(R) < self.archive['upper_size']:
            self.archive['set'] = copy.deepcopy(R)
            
            if len(self.archive['set']) != len(R):
                raise ValueError
            
            self.archive['size'] = len(self.archive['set'])
        else:
            sorted_R = crowding_distance_assignment(R)
            self.archive['set'].clear()
            for i in range(self.archive['upper_size']):
                self.archive['set'].append(sorted_R[i])
            if len(self.archive['set']) != self.archive['upper_size']:
                raise ValueError
            self.archive['size'] = len(self.archive['set'])    

            del sorted_R
        
        del S 
        del R   
        
    def update_swarms(self, ):

        ## Following the paper "Multiple Populations for Multiple Objectives: A Coevolutionary Technique for Solving Multiobjective Optimization Problems" Figure3
        ## Select an achieve solution randomly for each particle
        ## Update velocity and position using eq.6 and eq.7
        ## Eval
        for obj_index in range(self.OsDim):
            for i in range(self.config['NP']):
                self.swarms[obj_index][i].update_particle(self.archive, self.pBest, self.gBest)

        ## Update pBest
        self.update_pBest()

        ## Update gBest
        self.update_gBest()

        print('Current interiaw in CMPSO: {:.2f}'.format(self.config['inertiaw']))